from fitframework.api.decorators import fitable

class BasicQuery:
    question_dsl_examples = {}


class SingleIndicator(BasicQuery):
    question_dsl_examples = {
        "24年1月吕墅支行阳光贷申请客户数多少。时间：2024-01；吕墅支行：机构号为'吕墅支行'；指标编号为'阳光贷申请客户数'": {
            'data_date': '2024-01',
            'organ_code': '吕墅支行',
            'kpi_code': '阳光贷申请客户数'
        },
        "23年10月通江路支行阳光贷余额。时间：2023-10；通江路支行：机构号为'通江路支行'；指标编号为'阳光贷余额'": {
            'data_date': '2023-10',
            'organ_code': '通江路支行',
            'kpi_code': '阳光贷余额'
        }
    }


class SingleIndicatorMom(BasicQuery):
    question_dsl_examples = {
        "24年1月吕墅支行阳光贷申请客户数环比。时间：2024-01；吕墅支行：机构号为'吕墅支行'；指标编号为'阳光贷申请客户数'": {
            'data_date': '2024-01',
            'organ_code': '吕墅支行',
            'kpi_code': '阳光贷申请客户数'
        },
        "23年10月通江路支行阳光贷余额环比。时间：2023-10；通江路支行：机构号为'通江路支行'；指标编号为'阳光贷余额'": {
            'data_date': '2023-10',
            'organ_code': '通江路支行',
            'kpi_code': '阳光贷余额'
        }
    }


class SingleIndicatorPct(BasicQuery):
    question_dsl_examples = {
        "24年1月吕墅支行阳光贷申请客户数占比。时间：2024-01；吕墅支行：机构号为'吕墅支行'；指标编号为'阳光贷申请客户数'": {
            'condition1': {
                'data_date': '2024-01',
                'organ_code': '吕墅支行',
                'kpi_code': '阳光贷申请客户数'
            },
            'condition2': {
                'data_date': '2024-01',
                'kpi_code': '阳光贷申请客户数'
            }
        },
        "23年10月通江路支行阳光贷余额占比。时间：2023-10；通江路支行：机构号为'通江路支行'；指标编号为'阳光贷余额'": {
            'condition1': {
                'data_date': '2023-10',
                'organ_code': '通江路支行',
                'kpi_code': '阳光贷余额'
            },
            'condition2': {
                'data_date': '2023-10',
                'kpi_code': '阳光贷余额'
            }
        }
    }



@fitable("genericable.example_generate", "fitable.example_generate")
def example_generate(type: str):
    prompt = ""
    prompt += "<Example>"
    if type == "mom":
        prompt = handle_one_type(prompt, SingleIndicatorMom())
    elif type == "pct":
        prompt = handle_one_type(prompt, SingleIndicatorPct())
    else:
        prompt = handle_one_type(prompt, SingleIndicator())
    prompt += "</Example>"

    return prompt

def handle_one_type(prompt: str, query: BasicQuery):
    for item in query.question_dsl_examples.keys():
        prompt += "<ExampleItem>"
        prompt += "原问题和提取出的信息为："
        prompt += item
        prompt += "填充后返回的json结构为："
        prompt += str(query.question_dsl_examples[item])
        prompt += "</ExampleItem>"
    return prompt
